from nltk.corpus import wordnet as wn
from graphviz import Digraph

CONFIG = {
    "class_name": ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'],
    'visited': set(),  # point
    'edge_visited': set(),
    'layer': 7,
    'node_style': {
        'leaf': {'style': "filled", 'color': "#ff6666", 'fontcolor': "white", 'fontsize': "15"},
        'non_leaf': {'style': "filled", 'color': "#008080", 'fontcolor': "white", 'fontsize': "15"},
    },
    'edge_style': {
        'light': {'color': '#FF3030'},  # 突出
        'dark': {'color': '#838B8B'},
    },
    'light_root': None,
    'light_set': set(),
    'light_edge_visited': set(),
    'uuid': 0
}

digraph = Digraph("cifar10")


def update_uuid():
    CONFIG['uuid'] += 1


class Node():
    def __init__(self, name, setf):
        self.name = name
        self.setf = setf
        self.id = CONFIG['uuid']
        update_uuid()


def get_father(node):
    if not isinstance(node, list):
        pass
    nodes = []
    if not isinstance(node.setf, list):
        node.setf = [node.setf]
    for n in node.setf:
        fathers = n.hypernyms()
        for f in fathers:
            nodes.append(Node(lean_name(f), f))
    return nodes


def is_in_tree(node):
    return node.name in CONFIG['visited']


def lean_name(node):
    return str(node._name).split(".")[0]


def join_node(nodes):
    next_nodes = []
    for node in nodes:
        if not is_in_tree(node):
            digraph.node(name=node.name, **CONFIG['node_style']['leaf'])
            CONFIG['visited'].add(node.name)
        fathers = get_father(node)
        for father in fathers:
            if not is_in_tree(father):
                digraph.node(name=father.name, **CONFIG['node_style']['non_leaf'])
                CONFIG['visited'].add(father.name)
            next_nodes.append(father)

            cur_edge = "#".join([father.name, node.name])

            if node.name in CONFIG['light_set'] and cur_edge not in CONFIG['light_edge_visited']:  # 优先考虑是否突出显示
                digraph.edge(father.name, node.name, **CONFIG['edge_style']['light'])
                CONFIG['light_set'].add(father.name)
                CONFIG['light_edge_visited'].add(cur_edge)
                # next_nodes.append(father)
            elif cur_edge in CONFIG['edge_visited']:  # 是否 visited 的边
                continue
            else:
                digraph.edge(father.name, node.name, **CONFIG['edge_style']['dark'])
            CONFIG['edge_visited'].add(cur_edge)

    return next_nodes


def to_synset(nodes_name):
    res = []
    # for nn in nodes_name:
    #     words = wn.synsets(nn)
    #     for w in words:
    #         if lean_name(w) in CONFIG['class_name']:
    #             res.append(w)

    for i, nn in enumerate(nodes_name):
        res.append(Node(nn, wn.synsets(nn)))

    return res


def generate_tree():
    CONFIG['light_set'].add(CONFIG['light_root'])
    next_nodes = to_synset(CONFIG['class_name'])
    # next_nodes = CONFIG['class_name']
    for i in range(CONFIG['layer']):
        next_nodes = join_node(next_nodes)


if __name__ == "__main__":
    CONFIG['light_root'] = 'dog'
    generate_tree()
    digraph.view(directory="output/wordnet-cifar10/")
    # print(CONFIG['edge_visited'])